/*
    Author: @5mukx
*/

use std::ptr::null_mut;
use winapi::ctypes::c_void;
use winapi::shared::basetsd::ULONG_PTR;
use winapi::shared::ntdef::{HANDLE, UNICODE_STRING};
use winapi::um::winnt::{PROCESS_ALL_ACCESS, THREAD_ALL_ACCESS};

use ntapi::{
    ntrtl::{
        RtlAllocateHeap, RtlCreateProcessParametersEx, RtlDestroyProcessParameters, RtlFreeHeap,
        RtlInitUnicodeString, RtlProcessHeap, PRTL_USER_PROCESS_PARAMETERS, RTL_USER_PROC_PARAMS_NORMALIZED,
    },
    ntpsapi::{
        NtCreateUserProcess, PsCreateInitialState, PS_ATTRIBUTE_IMAGE_NAME, PS_ATTRIBUTE_LIST, PS_CREATE_INFO
    },
    ntpsapi::PS_ATTRIBUTE
};

use widestring::U16CString;

fn main() {
    unsafe {
        let mut nt_image_path: UNICODE_STRING = std::mem::zeroed();
        let path = U16CString::from_str(r"\??\C:\Windows\System32\calc.exe").unwrap();
        RtlInitUnicodeString(&mut nt_image_path, path.as_ptr() as *mut u16);

        // create process parameters
        let mut process_parameters: PRTL_USER_PROCESS_PARAMETERS = null_mut();
        let status = RtlCreateProcessParametersEx(
            &mut process_parameters,
            &mut nt_image_path,
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            RTL_USER_PROC_PARAMS_NORMALIZED,
        );
        
        if status < 0 {
            panic!("RtlCreateProcessParametersEx failed with status: {}", status);
        }

        println!("[i] Created Process parameter");


        let mut create_info: PS_CREATE_INFO = std::mem::zeroed();

        create_info.Size = std::mem::size_of::<PS_CREATE_INFO>();
        create_info.State = PsCreateInitialState;

        let attribute_list = RtlAllocateHeap(
            RtlProcessHeap(),
            0x00000008,// HEAP_ZERO_MEMORY,
            std::mem::size_of::<PS_ATTRIBUTE_LIST>(),
        ) as *mut PS_ATTRIBUTE_LIST;

        
        if attribute_list.is_null() {
            RtlDestroyProcessParameters(process_parameters);
            panic!("Failed to allocate attribute list");
        }

        println!("[*] Attribute List: {:?}", attribute_list);

        // Correct TotalLength
        (*attribute_list).TotalLength = std::mem::size_of::<PS_ATTRIBUTE_LIST>();

        // zero out Attributes[0] to prevent corruption
        std::ptr::write_bytes(&mut (*attribute_list).Attributes[0], 0, 1);

        // set PS_ATTRIBUTE
        (*attribute_list).Attributes[0] = PS_ATTRIBUTE {
            Attribute: PS_ATTRIBUTE_IMAGE_NAME,
            Size: nt_image_path.Length as usize,
            u: ntapi::ntpsapi::PS_ATTRIBUTE_u { Value: nt_image_path.Buffer as ULONG_PTR },
            ReturnLength: null_mut(),
        };

        // ensure handles are null before calling API
        let mut h_process: HANDLE = null_mut();
        let mut h_thread: HANDLE = null_mut();


        let status = NtCreateUserProcess(
            &mut h_process,
            &mut h_thread,
            PROCESS_ALL_ACCESS,
            THREAD_ALL_ACCESS,
            null_mut(),
            null_mut(),
            0,
            0,
            process_parameters as *mut c_void,
            &mut create_info,
            attribute_list,
        );

        RtlFreeHeap(RtlProcessHeap(), 0, attribute_list as *mut _);
        RtlDestroyProcessParameters(process_parameters);

        if status < 0 {
            panic!("NtCreateUserProcess failed with status: {:#x}", status);
        }

        println!("[i] NtCreateUserProcess Success");
    }
}
